package cn.jbolt.ai.core.chain.node.provider;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.config.AppConfigConstant;
import cn.jbolt.ai.app.entity.AiAppMessage;
import cn.jbolt.ai.app.service.AiAppMessageService;
import cn.jbolt.ai.app.service.AiAppService;
import cn.jbolt.ai.app.service.AiAppSessionService;
import cn.jbolt.ai.core.chain.ChainContext;
import cn.jbolt.ai.app.config.ChatNodeConfig;
import com.fasterxml.jackson.databind.JsonNode;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.model.JBoltAI;
import com.jboltai.event.EventState;
import com.jboltai.event.ai.aichat.AIChatEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import java.util.List;


@Component
public class ChatNodeProvider implements EventNodeProvider<AIChatEvent> {

    @Autowired
    private AiAppService aiAppService;
    @Autowired
    private AiAppSessionService aiAppSessionService;
    @Autowired
    private AiAppMessageService aiAppMessageService;
    private static final Logger logger = LoggerFactory.getLogger(ChatNodeProvider.class);

    public AIChatEvent provide(JsonNode nodeConfig, ChainContext context) {
        AIChatEvent event = JBoltAI.chat();
        try {
            //根据配置初始化event
            ChatNodeConfig config = configEvent(nodeConfig, context, event);
            event.onSuccess((e, res) -> {
                if (config.isOutputToResponse()) {
                    //如果配置项要求 内容返回到响应里，那就输出，默认是true的
                    //构建相应信息
                    context.sendResponseMsg(res.getContent(), config.getFormat());
                }
            });
            event.onFail((e, error) -> {
                if (config.isOutputToResponse()) {
                    //如果配置项要求 内容返回到响应里，那就输出
                    //构建相应信息
                    context.sendErrorMsg(error.getMsg());

                }
            });

            event.onComplete((e) -> {
                if (e.isSuccess()) { //失败
                    //事件执行成功
                    if (StrUtil.isNotBlank(config.getOutputField())) {
                        //配置项要求，把结果输出到上下文
                        context.setAttr(config.getOutputField(), e.getResult());
                    }
                    if (config.isPersistDb()) {//需要持久化数据
                        //把答案给存储
                        context.getResponseMsg().setContent(e.getResult().getContent());
                        context.getResponseMsg().setStatus(AiAppMessage.STATUS_SUCCESS);
                        aiAppMessageService.update(context.getResponseMsg());
                    }
                } else { //失败
                    if (config.isPersistDb()) {//需要持久化数据
                        context.getResponseMsg().setStatus(AiAppMessage.STATUS_ERROR);
                        if (e.getState() == EventState.FAIL) {
                            context.getResponseMsg().setErrorMsg(e.getError().getMsg());
                        } else if (e.getState() == EventState.CANCEL) {
                            context.getResponseMsg().setErrorMsg("响应被终止");
                        }
                        aiAppMessageService.update(context.getResponseMsg());
                    }

                }
            });
            return event;
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("ChatNodeHandler process 异常,异常原因", e);
            context.sendErrorMsg( "网络异常,请稍后重试....");
        }
        return null;
    }

    private ChatNodeConfig configEvent(JsonNode nodeConfig, ChainContext context, AIChatEvent event) {

        ChatNodeConfig config = ChatNodeConfig.from(nodeConfig.get(AppConfigConstant.CONFIG_TYPE_LLM));
        String prompt = generatePrompt(config, context);
        event.prompt(prompt);
        event.setStream(config.isStream());
        if (config.getModelType() != null) {
            event.setModel(config.getModelType());
        }
        if (StrUtil.isNotBlank(config.getDefaultModel())) {
            event.setModelName(config.getDefaultModel());
        }
        if (config.getMaxTokens() > 0) {
            event.setMaxTokens(config.getMaxTokens());
        }
        if (StrUtil.isNotBlank(config.getSystemRole())) {
            event.systemMessage(config.getSystemRole());
        }
        if (config.isUseContext()) {
            Long appId = context.getAiApp().getId();
            Long sessionId = context.getAiSession().getId();
            List<AIMessage> messageList = aiAppService.getAppSessionContextMessages(sessionId, appId);
            if (messageList != null && messageList.size() > 0) {
                event.contextMessages(messageList);
            }
        }
        if (config.getTemperature() != null&& config.getTemperature()>0) {
            event.setTemperature(config.getTemperature());
        }
        if (config.getTopP() != null&& config.getTopP()>0) {
            event.setTopP(config.getTopP());
        }
        if (config.getTopK() != null&& config.getTopK()>0) {
            event.setTopK(config.getTopK());
        }
        return config;

    }

    /**
     * 生成提示词
     *
     * @param config  配置
     * @param context 上下文
     * @return 处理后的提示词
     */
    private String generatePrompt(ChatNodeConfig config, ChainContext context) {
        if (StrUtil.isBlank(config.getPromptTpl())) {
            return context.getOriPrompt();
        }

        String template = config.getPromptTpl();
        if (context.getArgs() == null || context.getArgs().isEmpty()) {
            return template;
        }

        for (String key : context.getArgs().keySet()) {
            Object value = context.getArgs().get(key);
            if (value != null) {
                // 处理可能存在的空格情况，如 ${ key } 或 ${key}
                template = template.replaceAll("\\$\\{\\s*" + key + "\\s*\\}", value.toString());
            }
        }

        return template;
    }


}
